import coremltools as ct
from transformers import TFBertForSequenceClassification
import numpy as np
import os

CURRENT_PATH = os.path.dirname(os.path.abspath(__file__))

MODEL_SAVE_PATH = CURRENT_PATH + "/models/saved_model/"
CONVERTED_MODEL_PATH = CURRENT_PATH + "/models/converted_model/bert_model.mlmodel"

def convert_to_coreml():
    model = TFBertForSequenceClassification.from_pretrained(MODEL_SAVE_PATH, from_pt=False)

    # 打印模型输入信息
    model.summary()  # 打印模型摘要

    # 定义输入规格
    input_spec = ct.TensorType(name="input_ids", shape=(1, 50), dtype=np.int32)  # 使用 np.int32
    input_spec2 = ct.TensorType(name="attention_mask", shape=(1, 50), dtype=np.int32)  # 确保使用支持的类型
    input_spec3 = ct.TensorType(name="token_type_ids", shape=(1, 50), dtype=np.int32)  # 根据需要添加

    # 转换模型
    converted_model = ct.convert(
        model,
        inputs=[input_spec, input_spec2, input_spec3],
        convert_to="neuralnetwork"  # 指定转换为 neuralnetwork 类型
    )

    # 保存转换后的模型
    converted_model.save(CONVERTED_MODEL_PATH)
    print(f"Core ML model saved to {CONVERTED_MODEL_PATH}")

if __name__ == "__main__":
    convert_to_coreml()
